import pandas as pd
from app_config import get_engine_ts
import numpy as np


def str6(value):
    # 将数字转换为字符串
    str_value = str(value)
    # 如果字符串长度不足6位，则在前面补0
    formatted_value = str_value.zfill(6)
    return formatted_value


def index_weight(trade_date, from_file, to_file, head_key):

    # pre_000300_index = get_pro().index_weight(index_code='000300.SH', start_date=trade_date,
    #                                           end_date=trade_date)
    index_000016_index = pd.read_excel(from_file)
    index_000016_index['symbol'] = index_000016_index[head_key].apply(str6)

    query = f""" select symbol, ts_code from stock_basic """
    stock_basic = pd.read_sql(query, get_engine_ts())

    index_000016_index = pd.merge(index_000016_index, stock_basic, on='symbol', how='left')
    index_000016_index['con_code'] = index_000016_index['ts_code']
    sz50_ts_code = index_000016_index['con_code'].tolist()

    # 构建查询语句
    query = f"""
           SELECT ts_code,close,total_share,free_share FROM `daily_basic`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in sz50_ts_code)})
           AND trade_date = '{trade_date}'
           """
    # 执行查询并将结果转换为DataFrame
    stock_share = pd.read_sql_query(query, get_engine_ts())
    sz50 = pd.merge(index_000016_index, stock_share, left_on="con_code", right_on="ts_code", how='left')

    sz50['calc_percent'] = sz50['free_share'] / sz50['total_share'] * 100

    # 定义生成 b 列的规则
    def generate_b(value):
        if value <= 15:
            return np.ceil(value)  # 上调至最接近的整数
        elif 15 < value <= 20:
            return 20
        elif 20 < value <= 30:
            return 30
        elif 30 < value <= 40:
            return 40
        elif 40 < value <= 50:
            return 50
        elif 50 < value <= 60:
            return 60
        elif 60 < value <= 70:
            return 70
        elif 70 < value <= 80:
            return 80
        else:
            return 100

    sz50['adjust_percent'] = sz50['calc_percent'].apply(generate_b)
    sz50['adjust_share'] = sz50['adjust_percent'] / 100 * sz50['total_share']
    sz50['adjust_amount'] = sz50['adjust_share'] * sz50['close']

    amount_sum = sz50['adjust_amount'].sum()
    sz50['amount_precent'] = sz50['adjust_amount'] / amount_sum * 100

    sz50.to_excel(to_file)


if __name__ == '__main__':
    # index_weight()

    _trade_date = "20241202"
    _from_file = "file/000016cons.xls"
    _to_file = "file/to_000016上证50.xlsx"
    _head_key = '成份券代码Constituent Code'
    index_weight(trade_date=_trade_date, from_file=_from_file, to_file=_to_file, head_key=_head_key)

    _trade_date = "20241202"
    _from_file = "file/000016cons_1.xls"
    _to_file = "file/to_000016上证50_1.xlsx"
    _head_key = '成分券代码Constituent Code'
    index_weight(trade_date=_trade_date, from_file=_from_file, to_file=_to_file, head_key=_head_key)

